import { Knex } from 'knex';
import { BaseRepository, IBasicRepository } from './baseRepository';
import {
  camelCase,
  isPlainObject,
  mapKeys,
  mapValues,
  snakeCase,
  isEmpty,
} from 'lodash';
import { DataSourceName } from '@server/types';
import {
  IbisRedshiftConnectionType,
  IbisDatabricksConnectionType,
} from '@server/adaptors/ibisAdaptor';

export interface BIG_QUERY_CONNECTION_INFO {
  projectId: string;
  datasetId: string;
  credentials: string;
}
export interface POSTGRES_CONNECTION_INFO {
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  ssl: boolean;
}

export interface MYSQL_CONNECTION_INFO {
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  ssl: boolean;
}

export interface ORACLE_CONNECTION_INFO {
  user: string;
  password: string;
  host?: string;
  port?: number;
  database?: string;
  dsn?: string;
}

export interface MS_SQL_CONNECTION_INFO {
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  trustServerCertificate: boolean;
}

export interface CLICK_HOUSE_CONNECTION_INFO {
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  ssl: boolean;
}

export interface TRINO_CONNECTION_INFO {
  host: string;
  port: number;
  schemas: string;
  username: string;
  password: string;
  ssl: boolean;
}

export interface SNOWFLAKE_CONNECTION_INFO {
  user: string;
  account: string;
  database: string;
  schema: string;
  password?: string;
  privateKey?: string;
  warehouse?: string;
}

export interface DUCKDB_CONNECTION_INFO {
  initSql: string;
  extensions: Array<string>;
  configurations: Record<string, any>;
}

export interface ATHENA_CONNECTION_INFO {
  database: string;
  schema: string;
  s3StagingDir: string;
  awsRegion: string;
  awsAccessKey?: string;
  awsSecretKey?: string;
  webIdentityToken?: string;
  roleArn?: string;
  roleSessionName?: string;
}

export interface REDSHIFT_PASSWORD_AUTH {
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  redshiftType: IbisRedshiftConnectionType;
}

export interface REDSHIFT_IAM_AUTH {
  clusterIdentifier: string;
  user: string;
  database: string;
  awsRegion: string;
  awsAccessKey: string;
  awsSecretKey: string;
  redshiftType: IbisRedshiftConnectionType;
}

export interface DATABRICKS_PERSONAL_ACCESS_TOKEN_AUTH {
  serverHostname: string;
  httpPath: string;
  accessToken: string;
  databricksType: IbisDatabricksConnectionType;
}

export interface DATABRICKS_SERVICE_PRINCIPAL_AUTH {
  serverHostname: string;
  httpPath: string;
  clientId: string;
  clientSecret: string;
  azureTenantId?: string;
  databricksType: IbisDatabricksConnectionType;
}

export type REDSHIFT_CONNECTION_INFO =
  | REDSHIFT_PASSWORD_AUTH
  | REDSHIFT_IAM_AUTH;

export type DATABRICKS_CONNECTION_INFO =
  | DATABRICKS_PERSONAL_ACCESS_TOKEN_AUTH
  | DATABRICKS_SERVICE_PRINCIPAL_AUTH;

export type WREN_AI_CONNECTION_INFO =
  | ATHENA_CONNECTION_INFO
  | BIG_QUERY_CONNECTION_INFO
  | POSTGRES_CONNECTION_INFO
  | MYSQL_CONNECTION_INFO
  | ORACLE_CONNECTION_INFO
  | DUCKDB_CONNECTION_INFO
  | MS_SQL_CONNECTION_INFO
  | CLICK_HOUSE_CONNECTION_INFO
  | TRINO_CONNECTION_INFO
  | SNOWFLAKE_CONNECTION_INFO
  | REDSHIFT_CONNECTION_INFO
  | DATABRICKS_CONNECTION_INFO;

export interface RecommendationQuestionResult {
  question: string;
  category: string; // category for the question
  sql: string; // validated sql for this question, can be used in generateAskDetail
}

export interface Project {
  id: number; // ID
  type: DataSourceName; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc
  version: string; // Project datasource version
  displayName: string; // Project display name
  catalog: string; // Catalog name
  schema: string; // Schema name
  sampleDataset: string; // Sample dataset name
  connectionInfo: WREN_AI_CONNECTION_INFO;
  language?: string; // Project language

  // The recommended questions generated by AI
  queryId?: string;
  questions?: RecommendationQuestionResult[];
  questionsStatus?: string;
  questionsError?: object;
}

export interface IProjectRepository extends IBasicRepository<Project> {
  getCurrentProject: () => Promise<Project>;
}

export class ProjectRepository
  extends BaseRepository<Project>
  implements IProjectRepository
{
  private jsonTypeColumns = ['questions', 'questions_error', 'connection_info'];

  constructor(knexPg: Knex) {
    super({ knexPg, tableName: 'project' });
  }

  public async getCurrentProject() {
    const projects = await this.findAll({
      order: 'id',
      limit: 1,
    });
    if (!projects.length) {
      throw new Error('No project found');
    }
    return projects[0];
  }

  public override transformFromDBData: (data: any) => Project = (data: any) => {
    if (!isPlainObject(data)) {
      throw new Error('Unexpected db data');
    }
    const formattedData = mapValues(data, (value, key) => {
      if (this.jsonTypeColumns.includes(key) && typeof value === 'string') {
        // should return {} if value is null / {}, use value ? {} : JSON.parse(value) will throw error when value is null
        return isEmpty(value) ? {} : JSON.parse(value);
      }
      if (key === 'type') {
        return DataSourceName[value];
      }
      return value;
    });
    const camelCaseData = mapKeys(formattedData, (_value, key) =>
      camelCase(key),
    );
    return camelCaseData as Project;
  };

  public override transformToDBData: (data: Project) => any = (
    data: Project,
  ) => {
    if (!isPlainObject(data)) {
      throw new Error('Unexpected db data');
    }
    const snakeCaseData = mapKeys(data, (_value, key) => snakeCase(key));
    const formattedData = mapValues(snakeCaseData, (value, key) => {
      if (this.jsonTypeColumns.includes(key) && typeof value !== 'string') {
        return JSON.stringify(value);
      }
      return value;
    });
    return formattedData;
  };
}
